package com.kapcb.async.executor;

import com.kapcb.async.callback.DefaultGroupCallback;
import com.kapcb.async.callback.IGroupCallback;
import com.kapcb.async.wrapper.WorkerWrapper;

import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.stream.Collectors;

/**
 * <a>Title: Async </a>
 * <a>Author: kapcb <a>
 * <a>Description: 类入口，可以根据自己情况调整 core 线程的数量 <a>
 *
 * @author Kpacb
 * @version 1.0
 * @date 2024/5/11 22:01
 * @since 1.0
 */
public class Async {

    /**
     * 默认不定长线程池
     */
    private static final ThreadPoolExecutor COMMON_POOL = (ThreadPoolExecutor) Executors.newCachedThreadPool();

    /**
     * 注意，这里是个 static，也就是只能有一个线程池。用户自定义线程池时，也只能定义一个
     */
    private static ExecutorService executorService;

    /**
     * 开始任务
     *
     * @param timeout
     * @param executorService
     * @param workerWrappers
     * @return
     * @throws ExecutionException
     * @throws InterruptedException
     */
    public static boolean beginWork(long timeout, ExecutorService executorService, List<WorkerWrapper<?, ?>> workerWrappers) throws ExecutionException, InterruptedException {
        if (workerWrappers == null || workerWrappers.isEmpty()) {
            return false;
        }

        // 保存线程池变量
        Async.executorService = executorService;
        //
        ConcurrentHashMap<String, WorkerWrapper<?, ?>> forParamUseWrappers = new ConcurrentHashMap<>();
        CompletableFuture[] futures = new CompletableFuture[workerWrappers.size()];
        for (int i = 0; i < workerWrappers.size(); i++) {
            WorkerWrapper<?, ?> workerWrapper = workerWrappers.get(i);
            futures[i] = CompletableFuture.runAsync(() -> workerWrapper.work(executorService, timeout, forParamUseWrappers), executorService);
        }
        try {
            CompletableFuture.allOf(futures).get(timeout, TimeUnit.MILLISECONDS);
            return true;
        } catch (TimeoutException exception) {
            Set<WorkerWrapper<?, ?>> set = new HashSet<>();
            totalWorkers(workerWrappers, set);
            for (WorkerWrapper<?, ?> workerWrapper : set) {
                workerWrapper.stopNow();
            }

            return false;
        }
    }

    public static boolean beginWork(long timeout, ExecutorService executorService, WorkerWrapper<?, ?>... workerWrapper) throws ExecutionException, InterruptedException {
        if (workerWrapper == null || workerWrapper.length == 0) {
            return false;
        }

        List<WorkerWrapper<?, ?>> workerWrappers = Arrays.stream(workerWrapper).collect(Collectors.toList());
        return beginWork(timeout, executorService, workerWrappers);
    }

    public static boolean beginWork(long timeout, WorkerWrapper<?, ?>... workerWrapper) throws ExecutionException, InterruptedException {
        return beginWork(timeout, COMMON_POOL, workerWrapper);
    }

    public static void beginWorkAsync(long timeout, IGroupCallback groupCallback, WorkerWrapper<?, ?>... workerWrapper) {
        beginWorkAsync(timeout, COMMON_POOL, groupCallback, workerWrapper);
    }

    /**
     * 异步执行,直到所有都完成,或失败后，发起回调
     *
     * @param timeout
     * @param executorService
     * @param groupCallback
     * @param workerWrapper
     */
    public static void beginWorkAsync(long timeout, ExecutorService executorService, IGroupCallback groupCallback, WorkerWrapper<?, ?>... workerWrapper) {
        if (groupCallback == null) {
            groupCallback = new DefaultGroupCallback();
        }

        IGroupCallback finalGroupCallback = groupCallback;
        if (executorService != null) {
            executorService.submit(() -> {
                try {
                    boolean success = beginWork(timeout, executorService, workerWrapper);
                    if (success) {
                        finalGroupCallback.success(Arrays.asList(workerWrapper));
                    } else {
                        finalGroupCallback.failure(Arrays.asList(workerWrapper), new TimeoutException());
                    }
                } catch (ExecutionException | InterruptedException e) {
                    e.printStackTrace();
                    finalGroupCallback.failure(Arrays.asList(workerWrapper), e);
                }
            });
        } else {
            COMMON_POOL.submit(() -> {
                try {
                    boolean success = beginWork(timeout, COMMON_POOL, workerWrapper);
                    if (success) {
                        finalGroupCallback.success(Arrays.asList(workerWrapper));
                    } else {
                        finalGroupCallback.failure(Arrays.asList(workerWrapper), new TimeoutException());
                    }
                } catch (ExecutionException | InterruptedException e) {
                    e.printStackTrace();
                    finalGroupCallback.failure(Arrays.asList(workerWrapper), e);
                }
            });
        }
    }


    /**
     * 总共多少个执行单元
     *
     * @param workerWrappers
     * @param set
     */
    @SuppressWarnings("unchecked")
    private static void totalWorkers(List<WorkerWrapper<?, ?>> workerWrappers, Set<WorkerWrapper<?, ?>> set) {
        set.addAll(workerWrappers);
        for (WorkerWrapper<?, ?> wrapper : workerWrappers) {
            if (wrapper.getNextWrappers() == null) {
                continue;
            }
            List<WorkerWrapper<?, ?>> wrappers = wrapper.getNextWrappers();
            totalWorkers(wrappers, set);
        }
    }


    /**
     * 关闭线程池
     */
    public static void shutdown() {
        shutdown(executorService);
    }

    /**
     * 关闭线程池
     *
     * @param executorService ExecutorService
     */
    public static void shutdown(ExecutorService executorService) {
        if (executorService != null) {
            executorService.shutdown();
        } else {
            COMMON_POOL.shutdown();
        }
    }

    public static String getThreadCount() {
        return "activeCount=" + COMMON_POOL.getActiveCount() +
                "  completedCount " + COMMON_POOL.getCompletedTaskCount() +
                "  largestCount " + COMMON_POOL.getLargestPoolSize();
    }

}